/*
 * @Author: dingyuwen ding_yuwen@163.com
 * @Date: 2023-02-28 10:45:03
 * @LastEditTime: 2023-09-04 09:39:00
 * @LastEditors: DingYuwen 43669762+DingYuwen@users.noreply.github.com
 * @Description:
 */
import { Collapse, Flex, useDisclosure } from '@chakra-ui/react';

import { createContext, useCallback, useState } from 'react';

import AddStudioToggle from './AddStudioToggle';
import ModelListBox from './Model';
import TrainStudio from './TrainStudio';

import IAISegmentedControl from '@/common/components/IAIMantineSegmentedControl';

interface TrainingContextType {
  modelChange: number;
  setModelChange?: (n: number) => void;
  handleOnAdd?: () => void;
}
export const TrainingContext = createContext<TrainingContextType>({
  modelChange: 0,
});

export type ModelType = 'base' | 'lora';

const trainModelControlData: { value: ModelType; label: string }[] = [
  {
    value: 'lora',
    label: 'LoRA模型',
  },
  {
    value: 'base',
    label: '基础大模型',
  },
];
export default function Training() {
  const { isOpen, onToggle, onClose } = useDisclosure();
  const [listType, setListType] = useState<ModelType>('lora');

  const [modelChange, setModelChange] = useState<number>(0);

  const handleOnAdd = useCallback(() => {
    onClose();
    document.scrollingElement && (document.scrollingElement.scrollTop = 0);
  }, [onClose]);

  return (
    <Flex
      width="full"
      flex="1"
      pt={10}
      flexDirection="column"
      marginX="auto"
      px={4}
    >
      <TrainingContext.Provider
        value={{ modelChange, setModelChange, handleOnAdd }}
      >
        <Flex justify="space-between" align="center">
          <IAISegmentedControl
            value={listType}
            data={trainModelControlData}
            onChange={(v: ModelType) => setListType(v)}
          />

          <AddStudioToggle isOpen={isOpen} onToggle={onToggle} />
        </Flex>
        <Collapse in={isOpen} animateOpacity>
          <TrainStudio type={listType} />
        </Collapse>

        <ModelListBox type={listType} />
      </TrainingContext.Provider>
    </Flex>
  );
}
